# ('HDF5:"G:\\test\\VNP21A1D\\VNP21A1D.A2012195.h25v04.001.2019041082825.h5"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/Emis_14', '[1200x1200] //HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/Emis_14 (8-bit unsigned character)')
# ('HDF5:"G:\\test\\VNP21A1D\\VNP21A1D.A2012195.h25v04.001.2019041082825.h5"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/Emis_15', '[1200x1200] //HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/Emis_15 (8-bit unsigned character)')
# ('HDF5:"G:\\test\\VNP21A1D\\VNP21A1D.A2012195.h25v04.001.2019041082825.h5"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/Emis_16', '[1200x1200] //HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/Emis_16 (8-bit unsigned character)')
# ('HDF5:"G:\\test\\VNP21A1D\\VNP21A1D.A2012195.h25v04.001.2019041082825.h5"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/LST_1KM', '[1200x1200] //HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/LST_1KM (16-bit unsigned integer)')
# ('HDF5:"G:\\test\\VNP21A1D\\VNP21A1D.A2012195.h25v04.001.2019041082825.h5"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/QC', '[1200x1200] //HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/QC (16-bit unsigned integer)')
# ('HDF5:"G:\\test\\VNP21A1D\\VNP21A1D.A2012195.h25v04.001.2019041082825.h5"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/View_Angle', '[1200x1200] //HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/View_Angle (8-bit unsigned character)')
# ('HDF5:"G:\\test\\VNP21A1D\\VNP21A1D.A2012195.h25v04.001.2019041082825.h5"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/View_Time', '[1200x1200] //HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/View_Time (8-bit unsigned character)')

import os

import h5py
from osgeo import gdal, gdalconst, osr


# 写入tif
def write_tif(file_path, data, geotransform, nodata, dataType):
    gdal_type = ''
    if dataType == 'int16':
        gdal_type = gdal.GDT_Int16
    elif dataType == 'int32':
        gdal_type = gdal.GDT_Int32
    elif dataType == 'float32':
        gdal_type = gdal.GDT_Float32
    elif dataType == 'float64':
        gdal_type = gdal.GDT_Float64
    elif dataType == 'byte':
        gdal_type = gdal.GDT_Byte
    elif dataType == 'uint16':
        gdal_type = gdal.GDT_UInt16
    elif dataType == 'uint32':
        gdal_type = gdal.GDT_UInt32
    # 相当于一个创建数据集的驱动
    driver = gdal.GetDriverByName("GTiff")
    # 根据数据的维度来确定行列数（图像大小）
    rows, cols = data.shape
    # 创建一个新的数据集，存储输出文件；1是波段
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    # 设置仿射变换矩阵
    dataset.SetGeoTransform(geotransform)
    # 设置投影
    projInfo = 'PROJCS["unnamed",GEOGCS["Unknown datum based upon the custom spheroid", DATUM["Not specified (based on custom spheroid)", SPHEROID["Custom spheroid",6371007.181,0]],PRIMEM["Greenwich",0], UNIT["degree",0.0174532925199433]], PROJECTION["Sinusoidal"],PARAMETER["longitude_of_center",0],PARAMETER["false_easting",0],PARAMETER["false_northing",0],UNIT["Meter",1]]', \
               'GEOGCS["Unknown datum based upon the Clarke 1866 ellipsoid", DATUM["Not specified (based on Clarke 1866 spheroid)", SPHEROID["Clarke 1866",6378206.4,294.9786982139006]], PRIMEM["Greenwich",0], UNIT["degree",0.0174532925199433]]'
    dataset.SetProjection(projInfo[0])
    # 1个波段
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    # 释放内存
    del dataset


def getVIIRSGeographicInfo(input_file):
    # Get info from the StructMetadata Object
    f_Metadata = input_file['HDFEOS INFORMATION']['StructMetadata.0'][()].split()
    # Info returned is of type Byte, must convert to string before using it
    f_Metadata_byte2str = [s.decode('utf-8') for s in f_Metadata]
    # Get upper left points
    ulc = [i for i in f_Metadata_byte2str if 'UpperLeftPointMtrs' in i]
    ulcLon = float(ulc[0].replace('=', ',').replace('(', '') \
                   .replace(')', '').split(',')[1])
    ulcLat = float(ulc[0].replace('=', ',').replace('(', '') \
                   .replace(')', '').split(',')[2])
    return ulcLon, 0, ulcLat, 0


# 获取VIIRS的所有数据集
def getDatasetList(input_file):
    all_h5_objs = []
    input_file.visit(all_h5_objs.append)
    all_datasets = [str(obj) for obj in all_h5_objs if \
                    isinstance(input_file[obj], h5py.Dataset) and 'GRIDS' in obj]
    return all_datasets


def get_viirs_transform(tile_path, nRow, nCol):
    tile_geot = None
    if 'VNP21A1D' in tile_path or 'VNP21A1N' in tile_path:
        f = h5py.File(tile_path, 'r')
        geoInfo = getVIIRSGeographicInfo(f)
        tile_geot = list(geoInfo)
        xRes, yRes = 0, 0
        # Cell size not specified in the metadata of VIIRS version 001
        if nRow == 1200:  # VIIRS VNP09A1, VNP09GA - 1km
            yRes = -926.6254330555555
            xRes = 926.6254330555555
        elif nRow == 2400:  # VIIRS VNP09H1, VNP09GA - 500m
            yRes = -463.31271652777775
            xRes = 463.31271652777775
        elif nRow == 3600 and nCol == 7200:  # VIIRS VNP09CMG
            yRes = -0.05
            xRes = 0.05
            # Set upper left dims for CMG product
            tile_geot[0] = -180.00
            tile_geot[2] = 90.00
        # Set cell size and data type for output files
        tile_geot.insert(1, xRes)
        tile_geot.insert(5, yRes)

    return tile_geot


# viirs转tif
def viirs2tif(viirs_path, tif_dir):
    projInfo = 'PROJCS["unnamed",GEOGCS["Unknown datum based upon the custom spheroid", DATUM["Not specified (based on custom spheroid)", SPHEROID["Custom spheroid",6371007.181,0]],PRIMEM["Greenwich",0], UNIT["degree",0.0174532925199433]], PROJECTION["Sinusoidal"],PARAMETER["longitude_of_center",0],PARAMETER["false_easting",0],PARAMETER["false_northing",0],UNIT["Meter",1]]', \
               'GEOGCS["Unknown datum based upon the Clarke 1866 ellipsoid", DATUM["Not specified (based on Clarke 1866 spheroid)", SPHEROID["Clarke 1866",6378206.4,294.9786982139006]], PRIMEM["Greenwich",0], UNIT["degree",0.0174532925199433]]'
    datasets_list = []
    if 'VNP21A1D' in viirs_path:
        subDatasetList = ['LST_1KM', 'QC']
        datasets_list.append(gdal.Open(r'HDF5:"{}"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/LST_1KM'.format(viirs_path)))
        datasets_list.append(gdal.Open(r'HDF5:"{}"://HDFEOS/GRIDS/VIIRS_Grid_Daily_1km_LST21/Data_Fields/QC'.format(viirs_path)))
        nRow, nCol = datasets_list[0].RasterYSize, datasets_list[0].RasterXSize
        # viirs仿射变换矩阵
        geotrans = get_viirs_transform(viirs_path, nRow, nCol)
        # print(geotrans)
        # viirs投影
        for i, dataset in enumerate(datasets_list):
            # todo nodata读不出来
            # nodata = dataset.GetRasterBand(1).GetNoDataValue()
            # 给viirs附上仿射变换矩阵和投影
            dataset.SetGeoTransform(geotrans)
            dataset.SetProjection(projInfo[0])
            data = dataset.GetRasterBand(1).ReadAsArray()
            # 写入tif
            if not os.path.exists(tif_dir + "/" + subDatasetList[i]):
                os.makedirs(tif_dir + "/" + subDatasetList[i])
            output_path = tif_dir + "/" + subDatasetList[i] + "/" + os.path.splitext(viirs_path.split("\\")[-1])[0] + "_{}.tif".format(subDatasetList[i])
            gdal.Warp(output_path, dataset, dstSRS=projInfo[0], srcNodata=0, dstNodata=0, format='GTiff', resampleAlg=gdalconst.GRIORA_NearestNeighbour)
            # write_tif(output_path, data, geotrans, 0, 'uint16')


def main(hdf_dir, tif_dir):
    basename = ".h5"
    hdf_ls = [os.path.join(hdf_dir, f) for f in os.listdir(hdf_dir) if f.endswith(basename)]
    for hdf_path in hdf_ls:
        # 按照区域划分输出文件夹
        new_tif_dir = ''
        if 'h25v04' in hdf_path:
            new_tif_dir = tif_dir + "/h25v04"
        elif 'h25v05' in hdf_path:
            new_tif_dir = tif_dir + "/h25v05"
        if not os.path.exists(new_tif_dir):
            os.makedirs(new_tif_dir)
        viirs2tif(hdf_path, new_tif_dir)
        print("{} done".format(hdf_path))


if __name__ == '__main__':
    viirs_dir = r"G:\test\VNP21A1D"
    tif_dir = r"G:\test\process_result\VNP21A1D"
    main(viirs_dir, tif_dir)
